import torch
import torch.nn as nn


class Swin_Unet(nn.Module):
    def __init__(self, p, backbone, pre_head, cross_task_decoder):
        super(Swin_Unet, self).__init__()
        self.tasks = p.TASKS.NAMES
        self.backbone = backbone
        self.pre_head = pre_head
        self.cross_task_decoder = cross_task_decoder

    def forward(self, x):
        x = self.backbone(x)
        task_feature = {t: self.pre_head[t](x) for t in self.tasks}
        out = self.cross_task_decoder(task_feature)

        return out


